package com.mjf

import ml.dmlc.xgboost4j.java.DMatrix.SparseType
import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, XGBoost}

/**
 * XGBoost Scala 示例
 */
object XGBoostScalaExample {
  def main(args: Array[String]): Unit = {

    val rowHeaders: Array[Long] = Array[Long](0, 2, 4, 7)
    val data: Array[Float] = Array[Float](1f, 2f, 4f, 3f, 3f, 1f, 2f)
    val colIndex: Array[Int] = Array[Int](0, 2, 0, 3, 0, 1, 2)
    val numColumn: Int = 4

    /**
     * CSP稀疏矩阵数据解析结果：
     * [[1f, 0f, 2f, 0f],
     *  [4f, 0f, 0f, 3f],
     *  [3f, 1f, 2f, 0f]]
     */
    val trainData: DMatrix = new DMatrix(rowHeaders, colIndex, data, SparseType.CSR, numColumn)
    trainData.setLabel(Array[Float](1f, 0f, 1f))

    val testData: DMatrix = new DMatrix(Array[Float](2f, 2f, 1f, 1f), 1, 4)

    // define parameters
    val paramMap: Map[String, Any] = List(
      "eta" -> 0.1,
      "max_depth" -> 2,
      "objective" -> "binary:logistic").toMap

    val watches: Map[String, DMatrix] = Map(("train", trainData), ("test", testData))

    // number of iterations
    val round: Int = 2

    // train the model
    val model: Booster = XGBoost.train(trainData, paramMap, round, watches)

    // run prediction
    val predTrain: Array[Array[Float]] = model.predict(testData)

    // save model to the file.
    model.saveModel("D:\\coding\\idea\\xgboost-stu\\output")

  }
}
